#===========================================================#
#  Test the non-regularized multinomial model 
#===========================================================#

# 1) simulate data  
# generate data
set.seed(11)
tmp <- rmm(choice ~x2+x3|x4|x1, data = da, doFit = FALSE)
beta <- rnorm(ncol(tmp@X), sd = 1)
da$choice <- simulateY(tmp)

f1 <- mlogit(choice ~ x2+x3|x4|x1, data = da)
f2 <- rmm(choice ~x2+x3|x4|x1, data = da)
f3 <- rmm(choice ~x2+x3|x4|x1, data = da, 
          control = rmm.control(newton = FALSE))
f4 <- rmm(choice ~x2+x3|x4|x1, data = da, 
          control = rmm.control(standardize=FALSE))
head(as.numeric(f1$coefficients))
head(as.numeric(coef(f2)))
head(as.numeric(coef(f3)))
head(as.numeric(coef(f4)))
head(beta)

V <- solve(f2@hess)  
sc <- 1/attr(f2@X, "scale")
V <- diag(sc) %*% V %*% diag(sc)
head(sqrt(diag(V)))
head(sqrt(as.numeric(diag(vcov(f1)))))

X <- f2@X
mu <- f2@mu
a <- t(X) %*% diag(mu) %*% X

# 2) Use the Fishing data
data("Fishing", package = "mlogit")
Fishing$subject <- as.integer(cut(1:nrow(Fishing), 200))
data <- mlogit.data(Fishing, varying = c(2:9), shape = "wide", choice = "mode")
data <- transform(data, price = scale(price),
                  catch = scale(catch), 
                  income = scale(income))
data <- mlogit.data(data, choice = "mode", shape = "long", alt.var = "alt", id.var = "subject")
set.seed(11)
cset <- rep(1, nrow(data))
w <- which(data$alt == "boat" & data$mode == FALSE)
w <- w[as.logical(rbinom(length(w), 1, 0.3))]
cset[w] <- 0
data$cset <- cset

formula <- mode ~price + catch
formula <- mode ~1||price
formula <- mode ~price + catch|income

f1 <- mlogit(formula, data = data)
f2 <- rmm(formula, data = data)
f3 <- rmm(formula, data = data,
          control = rmm.control(newton = FALSE))
as.numeric(f1$coefficients)
as.numeric(coef(f2))
as.numeric(coef(f3))

# 3) model with pre-specified weights
set.seed(11)
wts <- rlnorm(max(data$chid))
wts <- wts[data$chid]
formula <- mode ~price + catch|income
f1 <- mlogit(formula, data = data, weights = wts)
f2 <- rmm(formula, data = data, weights = wts)
f3 <- rmm(formula, data = data, weights = wts,
          control = rmm.control(newton = FALSE))
as.numeric(f1$coefficients)
as.numeric(coef(f2))
as.numeric(coef(f3))

# 4) model with pre-specified cset
formula <- mode ~price + catch|income
f1 <- mlogit(formula, data = data, subset = cset == 1)
f2 <- rmm(formula, data = data, subset = cset == 1)
f3 <- rmm(formula, data = data, subset = cset == 1,
          control = rmm.control(newton = FALSE, epsilon = 1e-8))
as.numeric(f1$coefficients)
as.numeric(coef(f2))
as.numeric(coef(f3))
#cset <- lapply(1:max(data$chid), function(x) as.logical(cset[data$chid == x]))
#f4 <- mlogit(formula, data, cset = cset, scales = F)
#f2@beta
#f4$estimate

# 5) simulate data using the attributes decomp
fm <- choice ~0 + factor(z1) +z2 + x1 + x2 + x3
tmp <- rmm(fm, data = da, doFit = FALSE)
beta <- rnorm(ncol(tmp@X), sd = 1)
da$choice <- simulateY(tmp)

f1 <- mlogit(fm, data = da)
f2 <- rmm(fm, data = da)
f3 <- rmm(fm, data = da,
          control = rmm.control(newton = FALSE))
as.numeric(f1$coefficients)
as.numeric(coef(f2))
as.numeric(coef(f3))

V <- solve(f2@hess)  
head(sqrt(diag(V)))

# 6) simulate mixture data
S <- 2
tmp <- rmm(mode ~price + catch, data = data, doFit = FALSE)
pi <- 0.3
dims <- tmp@dims
beta1 <- c(0.5, 0.5, -0.1,  -1, 0.5)#rnorm(dims["nB"])
beta2 <- c(-1, -0.5, 0, 1, -2)#rnorm(dims["nB"])
set.seed(13)
tmp@beta <- matrix(beta1)
choice1 <- simulateY(tmp)
tmp@beta <- matrix(beta2)
choice2 <- simulateY(tmp)  
choice <- choice1  
seg <- rbinom(nlevels(factor(data$subject)), 1, pi)
sid2 <- which(seg == 0)  
idx <- data$subject %in% sid2
choice[idx] <- choice2[idx]
data$mode  <- choice

f1 <- rmm(mode ~price + catch, data = data, ncluster = 2, 
          control = rmm.control(EM.iter= 500, newton=FALSE))
f2 <- rmm(mode ~price + catch, data = data, ncluster = 2, 
          control = rmm.control(EM.iter= 500, newton=TRUE))
coef(f1)
coef(f2)

#===========================================================#
#  Test the group Lasso multinomial model 
#===========================================================#


# 1) non-mixture model (Fader and Hardy type decomp)
# simulate data 
fm <- choice ~0 + z1 + (z2 + z3) + z4 + z5 + (x1 + x2) + x3
set.seed(11)
tmp <- rmm(fm, data = da, 
           doFit = FALSE, control = rmm.control(standardize=TRUE))
U <- attr(tmp@X, "U")
beta <- c(rep(0, 4), 
          1, -1, 1, -1, -1, 1, 
          0, -1, 1)
gvar <- tmp@gvar
for (i in 1:(length(gvar) - 1)){
  idx <- (gvar[i] + 1):gvar[i + 1]
  beta[idx] <- as.numeric(U[[i]] %*% beta[idx])
}
tmp@beta <- matrix(beta)
da$choice <- simulateY(tmp)

# mnl fit
f1 <- rmm(fm, data = da)
coef(f1)

# glasso 
f2 <- rmm(fm, data = da, nlambda = 10,
            control = rmm.control(lambda.min.ratio=0.05, trace = TRUE), model = "glmm")
as(f2@beta, "sparseMatrix") 


# 1) mixture model 
# case 1: the Fader Hardy model
# simulate data  
S = 2
pi <- 0.3
tmp <- rmm(choice ~0 + z1 + z2 + z3 + z4 + z5 + x1 + x2 + x3, data = da, 
           doFit = FALSE, control = rmm.control(standardize = TRUE))
beta1 <- c(rep(0, 6), 
           0.5, -1, -0.5, -0.5, -1, -0.5, -0.5, 0.5, -1,
           0, 1, 1)
beta2 <- c(-0.5, 0.5, -0.5, -0.5, 0.5, 1, -1, 0.5, -1,               
           rep(0, 6), 
           1, -1, 0)
set.seed(11)
tmp@beta <- matrix(beta1)
choice1 <- simulateY(tmp)
tmp@beta <- matrix(beta2)
choice2 <- simulateY(tmp)  
choice <- choice1
da$sid <- factor(da$sid)
seg <- rbinom(nlevels(da$sid), 1, pi)
sid2 <- which(seg == 0)  
idx <- da$sid %in% sid2
choice[idx] <- choice2[idx]
da$choice  <- choice

# mnl fit
f0 <- rmm(choice ~0 + z1 + z2 + z3 + z4 + z5 + x1 + x2 + x3, data = da, ncluster = 2,
          control = rmm.control(newton = TRUE, gamma = 0.001))
coef(f0)
head(f0@mu)

f1 <- rmm(choice ~0 + z1 + z2 + z3 + z4 + z5 + x1 + x2 + x3, data = da,
          ncluster = 2, model = "glmm", nlambda = 10, 
          control = rmm.control(trace = TRUE))
coef(f1, 7)
f1@pi

